""" OllyDbg 1.10 Disassembler/Assembler
    Coded by Lin Zishen """

import os.path
import re

from ctypes import *
from itertools import count

DWORD = c_ulong
WORD = c_ushort
BYTE = c_ubyte

TEXTLEN = 256                           # Maximal length of text string
NREGSTACK = 32                          # Length of stack trace buffer
MAXCMDSIZE = 16                         # Maximal length of 80x86 command

DISASM_SIZE = 0                         # Determine command size only
DISASM_DATA = 1                         # Determine size and analysis data
DISASM_TRACE = 2                        # Trace integer registers
DISASM_FILE = 3                         # Disassembly, no symbols/registers
DISASM_CODE = 4                         # Disassembly, registers undefined
DISASM_ALL = 5                          # Complete disassembly
DISASM_RTRACE = 6                       # Disassemble with run-trace registers

class t_asmmodel(Structure):
    """Model to search for assembler command"""
    _fields_ = [
        ('code', BYTE * MAXCMDSIZE),    # Binary code
        ('mask', BYTE * MAXCMDSIZE),    # Mask for binary code (0: bit ignored)
        ('length', DWORD),              # Length of code, bytes (0: empty)
        ('jmpsize', DWORD),             # Offset size if relative jump
        ('jmpoffset', DWORD),           # Offset relative to IP
        ('jmppos', DWORD),              # Position of jump offset in command
    ]

class t_operand(Structure):
    """Full decription of command's operand"""
    _fields_ = [
        ('optype', BYTE),               # DEC_xxx (mem) or DECR_xxx (reg,const)
        ('opsize', BYTE),               # Size of operand
        ('regscale', BYTE * 8),         # Scales of registers
        ('seg', BYTE),                  # Segment register
        ('opconst', DWORD),             # Constant
    ]

class t_disasm(Structure):
    """Results of disassembling"""
    _fields_ = [
        ('ip', DWORD),                  # Instrucion pointer
        ('dump', BYTE * TEXTLEN),       # Hexadecimal dump of the command
        ('result', BYTE * TEXTLEN),     # Disassembled command
        ('comment', BYTE * TEXTLEN),    # Brief comment
        ('opinfo', BYTE * TEXTLEN * 3),           # Comments to command's operands
        ('cmdtype', DWORD),             # One of C_xxx
        ('memtype', DWORD),             # Type of addressed variable in memory
        ('nprefix', DWORD),             # Number of prefixes
        ('indexed', DWORD),             # Address contains register(s)
        ('jmpconst', DWORD),            # Constant jump address
        ('jmptable', DWORD),            # Possible address of switch table
        ('adrconst', DWORD),            # Constant part of address
        ('immconst', DWORD),            # Immediate constant
        ('zeroconst', DWORD),           # Whether contains zero constant
        ('fixupoffset', DWORD),         # Possible offset of 32-bit fixups
        ('fixupsize', DWORD),           # Possible total size of fixups or 0
        ('jmpaddr', DWORD),             # Destination of jump/call/return
        ('condition', DWORD),           # 0xFF:unconditional, 0:false, 1:true
        ('error', DWORD),               # Error while disassembling command
        ('warnings', DWORD),            # Combination of DAW_xxx
        ('optype', DWORD * 3),          # Type of operand (extended set DEC_xxx)
        ('opsize', DWORD * 3),          # Size of operand, bytes
        ('opgood', DWORD * 3),          # Whether address and data valid
        ('opaddr', DWORD * 3),          # Address if memory, index if register
        ('opdata', DWORD * 3),          # Actual value (only integer operands)
        ('op', t_operand * 3),          # Full description of operand
        ('regdata', DWORD * 8),         # Registers after command is executed
        ('regstatus', DWORD * 8),       # Status of registers, one of RST_xxx
        ('addrdata', DWORD),            # Traced memory address
        ('addrstatus', DWORD),          # Status of addrdata, one of RST_xxx
        ('regstack', DWORD * NREGSTACK),# Stack tracing buffer
        ('rststatus', DWORD * NREGSTACK),# Status of stack items
        ('nregstack', DWORD),           # Number of items in stack trace buffer
        ('reserved', DWORD * 29),       # Reserved for plugin compatibility
    ]

imagebase = 0x00400000
syntax = 0x004D916C
lowercase = 0x004D9170
tabarguments = 0x004D9174
extraspace = 0x004D9178
putdefseg = 0x004D917C
showmemsize = 0x004D9180
shownear = 0x004D9184
shortstringcmds = 0x004D9188
sizesens = 0x004D9194
symbolic = 0x004D91A0
sizesensitive = 0x004D9194

def load_ollydbg():
    path = 'ollydbg.exx'
    try:
        path = os.path.join(os.path.dirname(__file__), path)
    except NameError:
        pass
    ollydbg = cdll.LoadLibrary(path)

    #
    # HACK: Patch some options
    #
    def relocate(ea):
        return cast(ea - imagebase + ollydbg._handle, POINTER(c_ulong))

    def write_dword(ea, value):
        value = c_ulong(value)
        memmove(relocate(ea), byref(value), sizeof(value))

    write_dword(lowercase, 1)
    write_dword(tabarguments, 1)
    write_dword(extraspace, 1)
    write_dword(shortstringcmds, 1)
    write_dword(sizesensitive, 1)  # 0=PUSHA/PUSHAD, 1=PUSHAW/PUSHAD, 2=PUSHAW/PUSHA
    return ollydbg

ollydbg = load_ollydbg()


def Disasm(src, srcsize, srcip, srcdec, disasm, disasmmode, threadid):
    """
    Disassembles command, determines its size and decodes operands.
    Returns size of the command.
    Disasm functionality depends on the selected mode and global disassembling/analysis options.
    See description of t_disasm for more details:

    src - pointer to binary command that must be disassembled;

    srcsize - size of src. Length of 80x86 commands is limited to MAXCMDSIZE bytes;

    srcip - address of the command;

    srcdec - pointer to decoding data produced by Analyzer or NULL if decoding data is absent. You must supply srcdec if you want to decode switch tables, constants and strings;

    disasm - pointer to t_disasm structure that receives results of disassembling;

    disasmmode - disassembly mode, one of DISASM_xxx. See desctiption of t_disasm and table above;
        DISASM_SIZE	    Fastest mode, only calculates command size
        DISASM_DATA	    Extracts most important data, no textual information
        DISASM_TRACE	Extracts most important data and traces contents of integer registers, no textual information
        DISASM_FILE	    Disassembles command in assumption that registers are undefined and symbolic names are invalid. Usually used to disassemble contents of file
        DISASM_CODE	    Disassembles command assuming that registers are undefined
        DISASM_ALL	    Complete and relatively slow disassembly

    threadid - identifier of thread containing registers, or NULL if registers are undefined.
    """
    return ollydbg._Disasm(src, srcsize, srcip, srcdec, disasm, disasmmode, threadid)

def disasm(data, ea=0):
    """
    Disassembles data. Returns (result, size) or None if any error.

    >>> disasm('33C0'.decode('hex'))
    ('xor     eax, eax', 2)

    >>> disasm(assemble('push 111'))
    ('push    111', 5)
    """
    disasm = t_disasm()
    n = Disasm(data, len(data), ea, 0, byref(disasm), DISASM_FILE, 0)
    if n <= 0:
        raise ValueError('disasm error:\n data: {data}\n ea: {ea:08X}\n returns: {n:08X}'.format(**locals()))
    return string_at(disasm.result), n

def Assemble(cmd, ip, model, attempt, constsize, errtext):
    """
    Function Assemble, as expected, converts command in ASCII form to binary 32-bit code.
    It shares command table with Disasm, so if some command can be disassembled,
    it can be assembled back too, with one exception: Assemble doesn't support 16-bit addresses.
    Some commands have more than one encoding.
    By calling Assemble with parameter attempt=0,1... and constsize=0,1,2,3 one can get alternative variants
    and then select the shortest possible form (this is how OllyDbg implements assembling).
    However, only one address form is generated in each case
    ([EAX*2] but not [EAX+EAX]; [EBX+EAX] but not [EAX+EBX]; [EAX] will not use SIB byte; no DS: prefix and so on).

    Assemble compiles imprecise commands (where, for example, R32 replaces any general-purpose 32-bit register).
    This allows to generate imprecise search patterns, where mask contains zeros at the position occupied in code by register).
    Returns number of bytes in assembled code or non-positive number in case of detected error or when variant selected by combination of attempt and constsize doesn't exist.
    This number is the negative position of error in the input command.

    cmd - pointer to zero-terminated ASCII command;

    ip - address of the generated binary code in memory;

    model - pointer to structure that receives machine code and mask;

    attempt - index of alternative verison of the command. Call Assemble with attempt=0,1,2... to obtain all possible versions of the command. Stop this sequence when Assemble reports error;

    constsize - requested size of address constant and immediate data. Call Assemble with constsize=0,1,2,3 to obtain all possible variants of the version selected by attempt;

    errtext - pointer to text buffer of length at least TEXTLEN that receives description of detected error.
    """
    return ollydbg._Assemble(cmd, ip, model, attempt, constsize, errtext)

def fixup_const(disasm):
    """ Add 0 prefix to constant strings to resolve
        the ambiguity between hex strings and identifiers. """

    def repl(m):
        return '0' + m.group(0)
    # For heaven's sake, all constants are uppercased and commands are lowercased. """
    return re.sub(r'\b[0-9A-F]+\b', repl, disasm)

def remove_short(disasm):
    """ Remove 'short' from jxx short somewhere instructions,
        since sometimes assembler is unable to use short form. """
    return re.sub(r'\bshort\s*', '', disasm)

def assemble(text, ea=0, getall=False):
    """ See Assemble() for details.
        If getall=True returns a set of tuples of (code, mask).

        >>> assemble('push dword ptr [ebx+0C]', 0x400000).encode('hex')
        'ff730c'

        >>> assemble('jmp 401000', 0x400000).encode('hex')
        'e9fb0f0000'

        >>> [(code.encode('hex'), mask.encode('hex')) for code, mask in assemble('mov eax,[const]', getall=True)]
        [('a100000000', 'ff00000000'), ('8b0500000000', 'ffff00000000')]

    """
    text = remove_short(fixup_const(text))
    model = t_asmmodel()
    errtext = create_string_buffer(TEXTLEN)
    results = set()
    for attempt in count():
        for constsize in [0, 1, 2, 3]:
            n = Assemble(text, ea, byref(model), attempt, constsize, errtext)
            if n <= 0:
                if results:
                    break
                errtext = string_at(errtext)
                raise ValueError('assemble failed:\n text: {text}\n ea: {ea:08X}\n returns: {n:08X}\n errtext: {errtext}\n'.format(**locals()))
            else:
                results.add((string_at(model.code, n), string_at(model.mask, n)))
        else:
            continue
        break
    if getall:
        return results
    else:
        return sorted(results, key=lambda x: len(x[0]))[0][0]

def match(pattern, binary, ea=0):
    """ Returns True if binary matched pattern otherwise returns None.

        >>> match('call const', assemble('call 401005', 0x401000))
        True
    """
    for code, mask in assemble(pattern, ea, getall=True):
        if len(code) != len(binary):
            continue
        if all(map(lambda b, c, m: ord(b) & ord(m) == ord(c), binary, code, mask)):
            return True
    return False


if __name__ == '__main__':
    import doctest
    doctest.testmod()

